﻿using System.Windows;
using System.Windows.Shapes;
using System.Windows.Media;
using System;
using System.Windows.Controls;
using System.Windows.Media.Imaging;
using System.IO;
using System.ComponentModel;
using System.Numerics;

namespace HydrogenAtom4 {
    /// <summary>
    /// MainWindow.xaml の相互作用ロジック
    /// </summary>
    public partial class MainWindow : Window {
        public MainWindow() {
            InitializeComponent();

            mBWOne.DoWork += BackgroundOutputHydrogenBoxel;
            mBWOne.RunWorkerCompleted += BackgroundOutputHydrogenBoxelCompleted;
            mBWAll.DoWork += BackgroundOutputAllHydrogenBoxel;
            mBWAll.RunWorkerCompleted += BackgroundOutputHydrogenBoxelCompleted;
        }

        int mImageWidth  = 512;
        int mImageHeight = 512;

        enum Axis {
            XY,
            XZ,
            YZ,
            XYZ
        };

        enum OrbitalType {
            ComplexOrbital,
            RealOrbital,
        };

        BackgroundWorker mBWOne = new BackgroundWorker();
        BackgroundWorker mBWAll = new BackgroundWorker();

        private float DispValue(Complex v, double probabilityScale, OrbitalType orbitalType) {
            switch (orbitalType) {
            case OrbitalType.ComplexOrbital:
                return (float)(v.Magnitude * v.Magnitude * probabilityScale);
            case OrbitalType.RealOrbital:
                return (float)(v.Real * v.Real * probabilityScale);
            default:
                System.Diagnostics.Debug.Assert(false);
                return 0;
            }
        }

        private void CalcMaxMin(Axis axis, int precision, int n, int l, int m, double axisScale, OrbitalType orbitalType, out double vMax, out double vMin) {
            vMax = Double.MinValue;
            vMin = Double.MaxValue;

            var h = new HydrogenAtom(l, m);

            switch (axis) {
            case Axis.XY:
                for (int i=0; i < precision*2; ++i) {
                    double x = ((double)i - precision) / precision * axisScale;
                    for (int j=0; j < precision*2; ++j) {
                        double y = ((double)j - precision) / precision * axisScale;
                        double z = 0.0;

                        var s = SphericalCoordinates.FromXYZ(x, y, z);

                        var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                        var vEval = DispValue(v, 1.0, orbitalType);
                        
                        if (vMax < vEval) {
                            vMax = vEval;
                        }
                        if (vEval < vMin) {
                            vMin = vEval;
                        }
                    }
                }
                break;
            case Axis.XZ:
                for (int i=0; i < precision*2; ++i) {
                    double x = ((double)i - precision) / precision * axisScale;
                    for (int k=0; k < precision*2; ++k) {
                        double z = ((double)k - precision) / precision * axisScale;
                        double y = 0.0;

                        var s = SphericalCoordinates.FromXYZ(x, y, z);

                        var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                        var vEval = DispValue(v, 1.0, orbitalType);

                        if (vMax < vEval) {
                            vMax = vEval;
                        }
                        if (vEval < vMin) {
                            vMin = vEval;
                        }
                    }
                }
                break;
            case Axis.YZ:
                for (int j = 0; j < precision * 2; ++j) {
                    double y = ((double)j - precision) / precision * axisScale;
                    for (int k = 0; k < precision * 2; ++k) {
                        double z = ((double)k - precision) / precision * axisScale;
                        double x = 0.0;

                        var s = SphericalCoordinates.FromXYZ(x, y, z);

                        var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                        var vEval = DispValue(v, 1.0, orbitalType);

                        if (vMax < vEval) {
                            vMax = vEval;
                        }
                        if (vEval < vMin) {
                            vMin = vEval;
                        }
                    }
                }
                break;
            case Axis.XYZ:
                for (int i = 0; i < precision * 2; ++i) {
                    double x = ((double)i - precision) / precision * axisScale;
                    for (int j=0; j < precision*2; ++j) {
                        double y = ((double)j - precision) / precision * axisScale;
                        for (int k = 0; k < precision * 2; ++k) {
                            double z = ((double)k - precision) / precision * axisScale;

                            var s = SphericalCoordinates.FromXYZ(x, y, z);

                            var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                            var vEval = DispValue(v, 1.0, orbitalType);

                            if (vMax < vEval) {
                                vMax = vEval;
                            }
                            if (vEval < vMin) {
                                vMin = vEval;
                            }
                        }
                    }
                }
                break;
            }
        }

        private void DrawCanvas(Axis axis, int n, int l, int m, double axisScale, double probabilityScale) {
            mCanvas.Children.Clear();
            var h = new HydrogenAtom(l, m);

            var im = new Image();

            var bm = new WriteableBitmap(
                mImageWidth,
                mImageHeight,
                96, 
                96, 
                PixelFormats.Gray32Float, 
                null);
            im.Source = bm;
            im.Stretch = Stretch.Uniform;
            im.HorizontalAlignment = HorizontalAlignment.Left;
            im.VerticalAlignment   = VerticalAlignment.Top;
            im.Width = mCanvas.Width;
            im.Height = mCanvas.Height;
            
            var px = new float[bm.PixelHeight*bm.PixelWidth];

            OrbitalType orbitalType = (OrbitalType)mListBoxOrbitalType.SelectedIndex;

            switch (axis) {
            case Axis.XY:
                for (int i=0; i < mImageWidth; ++i) {
                    double x = ((double)i - mImageWidth / 2) / (mImageWidth / 2) * axisScale;
                    for (int j=0; j < mImageHeight; ++j) {
                        double y = ((double)j - mImageHeight / 2) / (mImageHeight / 2) * axisScale;
                        double z = 0.0;

                        var s = SphericalCoordinates.FromXYZ(x, y, z);
                        var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                        var vDisp = DispValue(v, probabilityScale, orbitalType);

                        px[i + j * mImageWidth] = (float)vDisp;
                    }
                }
                break;
            case Axis.XZ:
                for (int i=0; i < mImageWidth; ++i) {
                    double x = ((double)i - mImageWidth / 2) / (mImageWidth / 2) * axisScale;
                    for (int j=0; j < mImageHeight; ++j) {
                        double z = ((double)j - mImageHeight / 2) / (mImageHeight / 2) * axisScale;
                        double y = 0.0;

                        var s = SphericalCoordinates.FromXYZ(x, y, z);

                        var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                        var vDisp = DispValue(v, probabilityScale, orbitalType);

                        px[i + j * mImageWidth] = (float)vDisp;
                    }
                }
                break;
            case Axis.YZ:
                for (int i = 0; i < mImageWidth; ++i) {
                    double y = ((double)i - mImageWidth / 2) / (mImageWidth / 2) * axisScale;
                    for (int j = 0; j < mImageHeight; ++j) {
                        double z = ((double)j - mImageHeight / 2) / (mImageHeight / 2) * axisScale;
                        double x = 0.0;

                        var s = SphericalCoordinates.FromXYZ(x, y, z);

                        var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                        var vDisp = DispValue(v, probabilityScale, orbitalType);

                        px[i + j * mImageWidth] = (float)vDisp;
                    }
                }
                break;
            default:
                System.Diagnostics.Debug.Assert(false);
                break;
            }
            bm.WritePixels(new Int32Rect(0, 0, bm.PixelWidth, bm.PixelHeight), px, bm.BackBufferStride, 0);
            bm.Freeze();

            mCanvas.Children.Add(im);
        }

        private void OutputToFile(int n, int l, int m, int resolution, double axisScale, double probabilityScale, OrbitalType orbitalType, BinaryWriter bw) {

            bw.Write(resolution);
            bw.Write(resolution);
            bw.Write(resolution);
            int nFrames = 1;
            bw.Write(nFrames);

            var h = new HydrogenAtom(l, m);

            for (int k = 0; k < resolution; ++k) {
                double z = ((double)k - resolution / 2) / (resolution / 2) * axisScale;
                for (int j = 0; j < resolution; ++j) {
                    double y = ((double)j - resolution / 2) / (resolution / 2) * axisScale;
                    for (int i = 0; i < resolution; ++i) {
                        double x = ((double)i - resolution / 2) / (resolution / 2) * axisScale;

                        var s = SphericalCoordinates.FromXYZ(x, y, z);
                        var v = h.WaveFunction(1, n, s.r, s.θ, s.φ);
                        var vDisp = DispValue(v, probabilityScale, orbitalType);

                        float outputValue = (float)vDisp;
                        bw.Write(outputValue);
                    }
                }
            }
        }

        private void UpdateGraph() {
            int n;
            Int32.TryParse(mTextBoxN.Text, out n);
            int l;
            Int32.TryParse(mTextBoxL.Text, out l);
            int m;
            Int32.TryParse(mTextBoxM.Text, out m);

            double axisScale;
            Double.TryParse(mTextBoxAxisScale.Text, out axisScale);

            var axis = (Axis)mListBoxAxis.SelectedIndex;
            var orbitalType = (OrbitalType)mListBoxOrbitalType.SelectedIndex;

            double vMax;
            double vMin;
            CalcMaxMin(axis, 32, n, l, m, axisScale, orbitalType, out vMax, out vMin);
            double absMax = Math.Abs(vMax) < Math.Abs(vMin) ? Math.Abs(vMin) : Math.Abs(vMax);
            Console.WriteLine("max={0}", absMax);
            double probabilityScale = 1.0 * (1.0 / absMax);

            DrawCanvas(axis, n, l, m, axisScale, probabilityScale);
        }

        private void OnUpdateClicked(object sender, RoutedEventArgs e) {
            UpdateGraph();
        }

        private void Window_Loaded(object sender, RoutedEventArgs e) {
            UpdateGraph();
        }

        private void ListBoxSelectionChanged(object sender, SelectionChangedEventArgs e) {
            if (mListBoxAxis == null || mListBoxOrbitalType == null) {
                return;
            }
            UpdateGraph();
        }

        class HydrogenParams {
            public int n;
            public int l;
            public int m;
            public int resolution;
            public double axisScale;
            public string outputPath;
            public OrbitalType orbitalType;
        }

        private void OnOutputClicked(object sender, RoutedEventArgs e) {
            var param = new HydrogenParams();
            Int32.TryParse(mTextBoxN.Text, out param.n);
            Int32.TryParse(mTextBoxL.Text, out param.l);
            Int32.TryParse(mTextBoxM.Text, out param.m);
            Int32.TryParse(mTextBoxResolution.Text, out param.resolution);
            Double.TryParse(mTextBoxAxisScale.Text, out param.axisScale);
            param.outputPath = mTextBoxOutputFileName.Text;
            param.orbitalType = (OrbitalType)mListBoxOrbitalType.SelectedIndex;

            mGroupBoxConfiguration.IsEnabled = false;
            mGroupBoxOutput.IsEnabled = false;

            mBWOne.RunWorkerAsync(param);
        }

        private void OnOutputAllClicked(object sender, RoutedEventArgs e) {
            var param = new HydrogenParams();
            Int32.TryParse(mTextBoxResolution.Text, out param.resolution);
            Double.TryParse(mTextBoxAxisScale.Text, out param.axisScale);
            param.outputPath = mTextBoxOutputFileName.Text;
            param.orbitalType = (OrbitalType)mListBoxOrbitalType.SelectedIndex;

            mGroupBoxConfiguration.IsEnabled = false;
            mGroupBoxOutput.IsEnabled = false;

            mBWAll.RunWorkerAsync(param);
        }

        private static string[] mAzimuthalQuantumLetters = {
            "s",
            "p",
            "d",
            "f",
            "g",
            "h",
            "i",
        };

        private void OutputHydrogenBoxel(HydrogenParams param) {
            double vMax;
            double vMin;
            CalcMaxMin(Axis.XYZ, 32, param.n, param.l, param.m, param.axisScale, param.orbitalType, out vMax, out vMin);
            double absMax = Math.Abs(vMax) < Math.Abs(vMin) ? Math.Abs(vMin) : Math.Abs(vMax);
            using (var bw = new BinaryWriter(File.Open(param.outputPath, FileMode.Create))) {
                OutputToFile(param.n, param.l, param.m, param.resolution, param.axisScale, 1.0 / absMax, param.orbitalType, bw);
            }
        }

        void BackgroundOutputHydrogenBoxel(object sender, DoWorkEventArgs e) {
            var param = e.Argument as HydrogenParams;
            OutputHydrogenBoxel(param);
        }

        void BackgroundOutputAllHydrogenBoxel(object sender, DoWorkEventArgs e) {
            var param = e.Argument as HydrogenParams;
            for (int n = 1; n <= 7; ++n) {
                for (int l = 0; l < n; ++l) {
                    for (int m = 0; m <= l; ++m) {
                        param.n = n;
                        param.l = l;
                        param.m = m;
                        string directory = System.IO.Path.GetDirectoryName(param.outputPath);

                        param.outputPath = string.Format("{0}\\{1}{2}{3}.bvox", directory, n, mAzimuthalQuantumLetters[l], m);
                        OutputHydrogenBoxel(param);
                    }
                }
            }
        }

        void BackgroundOutputHydrogenBoxelCompleted(object sender, RunWorkerCompletedEventArgs e) {
            mGroupBoxConfiguration.IsEnabled = true;
            mGroupBoxOutput.IsEnabled = true;
        }

    }
}
